import angr
import json
import os
from typing import Dict, List, Optional
import pickle

#Backend.add




    

class BigBinaryCfgBuilder:
    """
    简化的angr符号加载器
    专注于：加载符号数据 -> 验证符号 -> 创建angr项目 -> 注入符号
    最终输出：可用的angr.Project对象
    """
    
    def __init__(self, binary_path: str,caches_dir:str, symbols_json_path: str = None):
        """
        初始化加载器
        
        Args:
            binary_path: 二进制文件路径
            symbols_json_path: 符号表JSON文件路径，如果为None则自动查找
        """
        self.binary_path = binary_path
        
        # 自动查找符号文件
        if symbols_json_path is None:
            symbols_json_path = os.path.join(os.path.dirname(binary_path), 'symbols_export.json')
            if not os.path.exists(symbols_json_path):
                symbols_json_path = os.path.join(os.getcwd(), 'symbols_export.json')
        
        self.symbols_json_path = symbols_json_path
        self.symbols_data = None
        self.valid_symbols = []
        self.regions=[]
        
        self.name2iddr={}
        
        self.caches_dir=caches_dir
        self.proj=None
    
    def load_symbols_data(self) -> bool:
        """
        从JSON文件加载符号数据
        
        Returns:
            bool: 加载是否成功
        """
        try:
            if not os.path.exists(self.symbols_json_path):
                print(f"Symbols file not found: {self.symbols_json_path}")
                return False
            
            with open(self.symbols_json_path, 'r', encoding='utf-8') as f:
                self.symbols_data = json.load(f)
            
            print(f"Loaded symbols data: {self.symbols_data['symbol_count']} symbols")
            print(f"Binary: {self.binary_path}")
            
            return True
            
        except Exception as e:
            print(f"Error loading symbols: {e}")
            return False
    
    def validate_symbols(self) -> List[Dict]:
        """
        验证和过滤符号数据
        
        Returns:
            List[Dict]: 有效的符号列表
        """
        if not self.symbols_data:
            return []
        
        valid_symbols = []
        base_addr = self.symbols_data['binary_info']['base_address']
        
        print(f"Validating symbols (base: {hex(base_addr)})...")
        
        function_count = 0
        data_count = 0
        filtered_count = 0
        
        for symbol in self.symbols_data['symbols']:
            name = symbol['name']
            address = symbol['address']
            symbol_type = symbol['type']
            source = symbol.get('source', 'unknown')
            
            # 过滤无效地址
            if address == 0 or address < 0x1000:
                filtered_count += 1
                continue
            
            # 宽松的地址范围检查
            if address < base_addr or address > (base_addr + 0x50000000):
                if name != "unknown_name":
                    filtered_count += 1
                    continue
            
            # 确定符号类型
            is_function = (
                symbol_type == 0x00000500 or 
                source == 'ida_function' or 
                name.startswith('fixed_') or 
                'func' in name.lower()
            )
            
            if is_function:
                function_count += 1
            else:
                data_count += 1
            
            # 处理符号名称
            if name == "unknown_name" or not name:
                name = f"sym_{hex(address)}"
            
            valid_symbols.append({
                'name': name,
                'address': address,
                'is_function': is_function,
                'source': source
            })
        
        print(f"Valid symbols: {len(valid_symbols)} (functions: {function_count}, data: {data_count})")
        print(f"Filtered: {filtered_count}")
        
        return valid_symbols
    
    def create_project(self) -> Optional[angr.Project]:
        """
        创建angr项目
        
        Returns:
            Optional[angr.Project]: 创建的项目对象，失败返回None
        """
        try:
            print("Creating angr project...")
            
            if not os.path.exists(self.binary_path):
                print(f"Binary file not found: {self.binary_path}")
                return None
            
            base_addr = self.symbols_data['binary_info']['base_address']
            
            # 尝试不同的配置
            configs = [
                {
                    'main_opts': {
                        'backend': 'blob',
                        'arch': 'mips32',
                        'base_addr': base_addr,
                        'entry_point': base_addr + 0x1000
                    },
                    'auto_load_libs': False,
                    'load_debug_info': False
                },
                {
                    'main_opts': {
                        'backend': 'blob',
                        'arch': 'mips32',
                        'base_addr': base_addr
                    },
                    'auto_load_libs': False
                },
                {
                    'main_opts': {
                        'backend': 'blob',
                        'arch': 'mips32'
                    }
                }
            ]
            
            proj = None
            for i, config in enumerate(configs, 1):
                try:
                    print(f"  Trying config {i}...")
                    proj = angr.Project(self.binary_path, **config)
                    print(f"  Success with config {i}")
                    break
                except Exception as e:
                    print(f"  Config {i} failed: {e}")
                    continue
            
            if proj:
                print(f"Project created: {proj.arch}, entry: {hex(proj.entry)}")
                return proj
            else:
                print("All configurations failed")
                return None
                
        except Exception as e:
            print(f"Error creating project: {e}")
            return None
    
    def inject_symbols(self, proj: angr.Project) -> bool:
        """
        将符号注入到angr项目中
        
        Args:
            proj: angr项目对象
            
        Returns:
            bool: 注入是否成功
        """
        if not proj or not self.valid_symbols:
            print("Project or symbols not ready")
            return False
        
        print("Injecting symbols...")
        
        function_count = 0
        data_count = 0
        error_count = 0
        
        # 准备符号存储
        try:
            main_obj = proj.loader.main_object
            if not hasattr(main_obj, 'symbols_by_name'):
                main_obj.symbols_by_name = {}
            symbols_by_name = main_obj.symbols_by_name
        except:
            # 使用自定义存储
            if not hasattr(proj, '_custom_symbols'):
                proj._custom_symbols = {}
            symbols_by_name = proj._custom_symbols
        
        # 注入符号
        for symbol in self.valid_symbols:
            name = symbol['name']
            address = symbol['address']
            is_function = symbol['is_function']
            
            # 添加前缀避免冲突
            symbol_name = name if name.startswith(('fixed_', 'restored_')) else f"restored_{name}"
            
            try:
                if is_function:
                    # 创建或更新函数
                    if address not in proj.kb.functions:
                        from angr.knowledge_plugins.functions import Function
                        
                        symbol_name = symbol_name.replace('fixed_', '')
                        
                        func = Function(proj.kb.functions, address, name=symbol_name)
                        proj.kb.functions._add_node(address, func)
                        self.name2iddr[symbol_name] = address
                        #print(symbol_name, hex(address), "Function injected")
                    else:
                        proj.kb.functions[address].name = symbol_name
                        #print(symbol_name, hex(address), "Function fixed")
                    function_count += 1
                else:
                    data_count += 1
                
                # 添加到符号表
                symbols_by_name[symbol_name] = address
                
            except Exception as e:
                error_count += 1
                if error_count <= 5:  # 只显示前5个错误
                    print(f"  Error injecting {symbol_name}: {e}")
        
        print(f"Injection completed:")
        print(f"  Functions: {function_count}")
        print(f"  Data symbols: {data_count}")
        print(f"  Errors: {error_count}")
        
        return True
    
    def load_project(self) -> Optional[angr.Project]:
        """
        完整的加载流程：加载符号 -> 验证 -> 创建项目 -> 注入符号
        
        Returns:
            Optional[angr.Project]: 加载完成的项目对象
        """
        print("="*50)
        print("ANGR PROJECT LOADING")
        print("="*50)
        
        # 1. 加载符号数据
        if not self.load_symbols_data():
            return None
        
        # 2. 验证符号
        self.valid_symbols = self.validate_symbols()
        if not self.valid_symbols:
            print("No valid symbols found!")
            return None
        
        # 3. 创建项目
        proj = self.create_project()
        if not proj:
            return None
        
        # 4. 注入符号
        if not self.inject_symbols(proj):
            print("Warning: Symbol injection failed")
        
        print("="*50)
        print("PROJECT READY!")
        print(f"Functions: {len(proj.kb.functions)}")
        print(f"Architecture: {proj.arch}")
        print(f"Entry point: {hex(proj.entry)}")
        print("="*50)
        
        return proj

    def build(self,region_number=10,regions=None):
        
        if not self.proj:
            self.proj = self.load_project()
            if not self.proj:
                print("Failed to load project")
                return
        
        if regions is None:
            regions= BigBinaryCfgBuilder.default_split(self.proj,region_number)
        
        
        self.regions = regions
        
        
        BigBinaryCfgBuilder.build_big_cfg(
            proj=self.proj,
            caches_dir=self.caches_dir,
            regions=self.regions,
            region_names=[self.cfg_block_name(i) for i in range(len(self.regions))]    
        )
    
    def build_partial(self,select_indexes,region_number=10,regions=None):
        """
        构建指定索引的CFG块
        
        Args:
            select_indexes: 需要构建的索引列表
            region_number: 区域数量
        """
        if not self.proj:
            self.proj = self.load_project()
            if not self.proj:
                print("Failed to load project")
                return
            
        if regions is None:
            regions = BigBinaryCfgBuilder.default_split(self.proj, region_number)
        
        self.regions= regions
        
        for index in select_indexes:
            cfg_path = os.path.join(self.caches_dir, self.cfg_block_name(index))
            if not os.path.exists(cfg_path):
                print(f"Processing region {index}: {self.regions[index]}")
                cfg = self.proj.analyses.CFGFast(show_progressbar=True, regions=[self.regions[index]])
                with open(cfg_path, 'wb') as f:
                    pickle.dump(cfg, f)
                
                del cfg
    
    def load_partial_cfg(self,region_number=10,regions=None):
        """
        加载部分CFG块
        
        Args:
            region_number: 区域数量
        """
        if not self.proj:
            self.proj = self.load_project()
            if not self.proj:
                print("Failed to load project")
                return
        
        
        if regions is None:
            regions = BigBinaryCfgBuilder.default_split(self.proj, region_number)
        
        self.regions= regions
        
        for i, region in enumerate(self.regions):
            cfg_path = os.path.join(self.caches_dir, self.cfg_block_name(i))
            if os.path.exists(cfg_path):
                with open(cfg_path, 'rb') as f:
                    cfg = pickle.load(f)
                    yield cfg
                
                
        
    def cfg_block(self,index:int):
        """
        获取指定索引的CFG块
        :param index: 块索引
        :return: CFG块对象
        """
        
        assert len(self.regions) > 0, f"Index {index} out of range for regions"
        
        cfg_path = os.path.join(self.caches_dir, self.cfg_block_name(index))
        if not os.path.exists(cfg_path):
            print(f"CFG file not found: {cfg_path}")
            return None
        
        with open(cfg_path, 'rb') as f:
            cfg = pickle.load(f)
        
        return cfg
    
    def cfg_block_name(self,index:int):
        """
        获取指定索引的CFG块名称
        :param index: 块索引
        :return: CFG块名称
        """
        assert len(self.regions) > 0, f"Index {index} out of range for regions"
        
        start, end = self.regions[index]
        return f"region_{hex(start)}to{hex(end)}.cfg"
    
    def __getitem__(self, index: int):
        """
        支持索引访问，返回指定索引的CFG块
        :param index: 块索引
        :return: CFG块对象
        """
        return self.cfg_block(index)


    @staticmethod
    def build_big_cfg(proj:angr.Project,caches_dir:str,regions:list[tuple],region_names:list[str]=None):    
        if not proj:
            print("Failed to load project")
            return None
        
        assert len(regions)==len(region_names), "Regions and region names must have the same length"
        
        if not os.path.exists(caches_dir):
            os.makedirs(caches_dir)
            
        assert len(os.listdir(caches_dir)) == 0, "Cache directory must be empty"
        
        for i, region in enumerate(regions):
            #start,end= region
            region_path = os.path.join(caches_dir, region_names[i])
            if not os.path.exists(region_path):
                print(f"Processing region {i}: {region}")
                cfg = proj.analyses.CFGFast(show_progressbar=True,
                                            regions=[region])
                
                
                with open(region_path, 'wb') as f:
                    pickle.dump(cfg, f)
                
                del cfg  # 释放内存    
                print(f"Region {i} saved to {region_path}")
                
            else:
                print(f"Region {i} already exists, skipping")

    @staticmethod
    def default_split(proj:angr.Project,region_number):
        """
        将项目分割成指定数量的区域
        :param proj: angr.Project对象
        :param region_number: 区域数量
        :return: 分割后的区域列表
        """
        regions = []
        start_addr = proj.loader.main_object.min_addr
        end_addr = proj.loader.main_object.max_addr
        
        if region_number==1:
            # 如果只需要一个区域，直接返回整个范围
            return [(start_addr, end_addr)]
        
        region_size = (end_addr - start_addr) // region_number
        
        
        
        func_addrs= [func.addr for func in proj.kb.functions.values()]
        func_addrs.sort()
        
        func_index=0
        
        # 将最接近region_size的函数地址作为划分参考
        for i in range(region_number):
            region_start_pred = start_addr + i * region_size
            region_end_pred = start_addr + (i + 1) * region_size
            
            # 找到第一个函数地址大于region_start的函数
            for i in range(func_index, len(func_addrs)):
                if func_addrs[i] >= region_start_pred:
                    region_start_pred = func_addrs[i]
                    func_index = i
                    break
            
            # 找到最后一个函数地址小于region_end的函数
            for j in range(func_index, len(func_addrs)):
                if func_addrs[j] > region_end_pred:
                    region_end_pred = func_addrs[j-1]
                    func_index = j-1
                    break
                
            regions.append((region_start_pred, region_end_pred))
        
        return regions

    @property
    def base_address(self) -> int:
        """
        获取二进制文件的基地址
        
        Returns:
            int: 基地址
        """
        if self.symbols_data:
            return self.symbols_data['binary_info']['base_address']
        return 0x1000
    
    @property
    def func_start_addr(self) -> int:
        """
        获取函数起始地址
        
        Returns:
            int: 函数起始地址
        """
        if self.symbols_data:
            first_func = None
            for symbol in self.symbols_data['symbols']:
                if symbol['source']=='ida_function':
                    # 返回第一个函数的地址
                    first_func = symbol
                    break
            if first_func:
                return first_func['address']
        return 0x1000
    
    @property
    def func_end_addr(self) -> int:
        """
        获取函数结束地址
        
        Returns:
            int: 函数结束地址
        """
        if self.symbols_data:
            last_func = None
            for symbol in reversed(self.symbols_data['symbols']):
                if symbol['source']=='ida_function':
                    # 返回最后一个函数的地址
                    last_func= symbol
                    break
            if last_func:
                # 返回最后一个函数的地址
                return last_func['address']+last_func['size']
        return 0x50000000
    
# 便捷函数
def load_angr_project(binary_path: str, symbols_json_path: str = None) -> Optional[angr.Project]:
    """
    便捷函数：一步加载带符号的angr项目
    
    Args:
        binary_path: 二进制文件路径
        symbols_json_path: 符号文件路径（可选）
    
    Returns:
        Optional[angr.Project]: 加载的项目对象
    """
    loader = BigBinaryCfgBuilder(binary_path, symbols_json_path)
    return loader.load_project()

def static_cfg(binary_path:str,caches_path:str,regions:list[tuple],symbols_json_path:str=None):
    proj= load_angr_project(binary_path, symbols_json_path)
    cfg=None
    if os.path.exists(caches_path):
        data=pickle.load(open(caches_path, 'rb'))
        cfg=data
    
    else:
        cfg=proj.analyses.CFGFast(show_progressbar=True,regions=regions)
        with open(caches_path, 'wb') as f:
            pickle.dump(cfg, f)    
        
        
    return cfg





# 使用示例
if __name__ == "__main__":
    bbinary_builder= BigBinaryCfgBuilder(
        binary_path='caches_data/ine.txtbin',
        caches_dir='caches_data/nokia_cfg',
        symbols_json_path='caches_data/symbols_export.json'
    )
    
    